import multiprocessing
import torch
from config import wzzCredictConfig
from pytorch_lightning import Trainer
from model.wzzModel import WzzModel_LM
from data.credictData import CreateDataset_LM
from pytorch_lightning.loggers import CSVLogger
from pytorch_lightning.callbacks import ModelCheckpoint
from utils.wzzUtils import readCreditData

if __name__ == '__main__':
    # 防止Windows训练多线程报错的问题
    multiprocessing.freeze_support()

    # 加速配置
    torch.set_float32_matmul_precision(precision='high')

    # 加载模型
    model = WzzModel_LM()

    # 加载数据
    data_array = readCreditData(file_path=wzzCredictConfig.DATASET_PATH)
    data_model = CreateDataset_LM(data_array=data_array,
                                  batch_size=wzzCredictConfig.BATCH_SIZE,
                                  num_workers=wzzCredictConfig.NUM_WORKERS)

    # 设置训练输出日志
    logger = CSVLogger(save_dir=wzzCredictConfig.LOG_PATH, name='AE')

    # 设置参数的保存
    checkpoint_callback = ModelCheckpoint(
        monitor='loss',
        dirpath=wzzCredictConfig.LOG_PATH,
        filename='AE-best-weight',
        save_top_k=1,
        mode='min'
    )

    # 开始训练
    trainer = Trainer(max_epochs=wzzCredictConfig.EPOCH, accelerator='gpu', devices=1, logger=logger, precision='bf16',
                      callbacks=checkpoint_callback, default_root_dir=wzzCredictConfig.LOG_PATH)
    trainer.fit(model=model, datamodule=data_model)
